"""
ConnLLM速率限制工具
提供API调用的速率限制功能
"""
import time
import logging
import threading
from typing import Dict, Any, Optional, Callable, TypeVar, Union, List, Tuple

# 设置日志
logger = logging.getLogger("connllm.rate_limit")

# 类型变量，用于装饰器的泛型返回类型
F = TypeVar('F', bound=Callable[..., Any])


# 提供商默认速率限制（每分钟请求数）
DEFAULT_RATE_LIMITS = {
	# 远程提供商
	"anthropic": 45,       # Anthropic默认限制
	"openai": 60,          # OpenAI默认限制
	"openrouter": 60,      # OpenRouter默认限制
	"moonshot": 60,        # Moonshot默认限制
	"ppinfra": 60,         # 鹏程立方默认限制
	"volces": 60,          # Volces默认限制
	
	# 本地提供商通常无需限制，但设置一个合理值以防止资源耗尽
	"ollama": 120,
	"vllm": 120,
	
	# 默认限制
	"default": 60
}


class RateLimiter:
	"""速率限制器，限制API调用频率"""
	
	def __init__(self, rate: int, per: float = 60.0):
		"""
		初始化速率限制器
		
		Args:
			rate: 时间窗口内允许的最大调用次数
			per: 时间窗口大小(秒)
		"""
		self.rate = max(1, rate)  # 至少允许1个请求
		self.per = max(0.1, per)  # 时间窗口至少0.1秒
		self.allowance = self.rate  # 当前剩余配额
		self.last_check = time.time()  # 上次检查时间
		self.lock = threading.RLock()  # 线程锁，确保线程安全
		
	def __call__(self, func: F) -> F:
		"""
		装饰器，限制函数调用频率
		
		Args:
			func: 要限制的函数
			
		Returns:
			装饰后的函数
		"""
		def wrapper(*args, **kwargs):
			# 请求令牌
			self.request_token()
			return func(*args, **kwargs)
			
		# 保留原函数的元数据
		wrapper.__name__ = func.__name__
		wrapper.__doc__ = func.__doc__
		wrapper.__annotations__ = func.__annotations__
		
		return wrapper  # type: ignore
		
	def request_token(self, tokens: int = 1) -> None:
		"""
		请求令牌，如果没有足够令牌则等待
		
		Args:
			tokens: 请求的令牌数量
		"""
		with self.lock:
			current = time.time()
			time_passed = current - self.last_check
			self.last_check = current
			
			# 根据时间推移增加令牌
			self.allowance += time_passed * (self.rate / self.per)
			
			# 不允许超过最大令牌数
			if self.allowance > self.rate:
				self.allowance = self.rate
				
			# 如果剩余令牌不足
			if self.allowance < tokens:
				# 需要等待的时间
				wait_time = (tokens - self.allowance) * (self.per / self.rate)
				logger.debug(f"速率限制：等待{wait_time:.2f}秒")
				time.sleep(wait_time)
				self.allowance = 0
			else:
				# 扣除令牌
				self.allowance -= tokens


# 全局速率限制器字典
_rate_limiters: Dict[str, RateLimiter] = {}
_rate_limiters_lock = threading.RLock()


def get_rate_limiter(provider: str, rate: Optional[int] = None) -> RateLimiter:
	"""
	获取特定提供商的速率限制器
	
	Args:
		provider: 提供商名称
		rate: 可选的速率限制，如果未指定则使用默认值
		
	Returns:
		速率限制器实例
	"""
	with _rate_limiters_lock:
		if provider not in _rate_limiters:
			# 使用指定的速率或默认值
			limit_rate = rate or DEFAULT_RATE_LIMITS.get(
				provider, DEFAULT_RATE_LIMITS["default"]
			)
			_rate_limiters[provider] = RateLimiter(limit_rate)
			
		return _rate_limiters[provider]


def rate_limit(func: Optional[F] = None, *, provider: str = "default", rate: Optional[int] = None) -> Union[F, Callable[[F], F]]:
	"""
	速率限制装饰器，限制函数调用频率
	
	可以直接用作装饰器，或者带参数使用
	
	Args:
		func: 要装饰的函数
		provider: 提供商名称，用于选择速率限制
		rate: 可选的速率限制，覆盖默认值
		
	Returns:
		装饰器或装饰后的函数
	"""
	def decorator(f: F) -> F:
		limiter = get_rate_limiter(provider, rate)
		return limiter(f)
		
	# 直接用作装饰器
	if func is not None:
		return decorator(func)
		
	# 带参数使用
	return decorator


def provider_rate_limit(provider_type: str, model: Optional[str] = None) -> Callable[[F], F]:
	"""
	提供商特定的速率限制装饰器
	
	Args:
		provider_type: 提供商类型
		model: 可选的模型名称，某些提供商可能基于模型有不同限制
		
	Returns:
		装饰器函数
	"""
	# 特定模型的速率限制（未来可扩展）
	if model and False:  # 当前未实现，预留接口
		pass
		
	return rate_limit(provider=provider_type)


def set_rate_limit(provider: str, rate: int, per: float = 60.0) -> None:
	"""
	设置提供商的速率限制
	
	Args:
		provider: 提供商名称
		rate: 时间窗口内允许的最大调用次数
		per: 时间窗口大小(秒)
	"""
	with _rate_limiters_lock:
		_rate_limiters[provider] = RateLimiter(rate, per)
		
		
def get_rate_limits() -> Dict[str, int]:
	"""
	获取所有提供商的当前速率限制
	
	Returns:
		提供商到速率限制的映射字典
	"""
	limits = {}
	
	# 添加默认限制
	for provider, rate in DEFAULT_RATE_LIMITS.items():
		limits[provider] = rate
		
	# 添加自定义限制
	with _rate_limiters_lock:
		for provider, limiter in _rate_limiters.items():
			limits[provider] = limiter.rate
			
	return limits
